package com.webcat.service.impl;

import com.webcat.service.RagService;
import jakarta.annotation.Resource;
import org.apache.tika.Tika;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import java.io.*;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Service
public class RagServiceImpl implements RagService {


    //@Resource
    private VectorStore vectorStore;

    @Override
    public void loadDocument() {
        TokenTextSplitter textSplitter = new TokenTextSplitter();
        // 存储 Embedding
        List<Document> documents = List.of(new Document("你叫王军，你是一个程序员"));
        vectorStore.accept(textSplitter.split(documents));
    }

    @Override
    public void loadDocument(List<MultipartFile> files) {
        if (files == null || files.isEmpty()) {
            return;
        }
        List<String> textList = new ArrayList<>();
        files.parallelStream().forEach(file -> {
            try(InputStream in = file.getInputStream()) {
                Tika tika = new Tika();
                String text = tika.parseToString(in);
                textList.add(text);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
        TokenTextSplitter textSplitter = new TokenTextSplitter();
        List<Document> documents = textList.stream().map(Document::new).collect(Collectors.toList());
        vectorStore.accept(textSplitter.split(documents));
    }

    @Override
    public void loadDocument(InputStream inputStream) {
        Tika tika = new Tika();
        String text;
        try {
            text = tika.parseToString(inputStream);
        }
        catch (Exception e) {
            throw new RuntimeException(e);
        }
        this.loadDocument(text);
    }

    @Override
    public void loadDocument(String text) {
        TokenTextSplitter textSplitter = new TokenTextSplitter();
        // 存储 Embedding
        List<Document> documents = List.of(new Document(text));
        vectorStore.accept(textSplitter.split(documents));
    }

    @Override
    public List<Document> query(String text) {
        return vectorStore.similaritySearch(text);
    }
}
